package com.chatplus.application.listener;

import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.chatplus.application.aiprocessor.platform.image.sd.SdImageProcessor;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.common.util.OkHttpClientUtil;
import com.chatplus.application.domain.dto.ApiKeyDto;
import com.chatplus.application.domain.entity.draw.SdJobEntity;
import com.chatplus.application.service.draw.SdJobService;
import okhttp3.Call;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.Response;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Component;

import java.io.IOException;

/**
 * SD 进度条处理
 *
 * @Author Angus
 * @Date 2024/3/30
 */
@Component
public class SdJobProcess {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(SdJobProcess.class);

    private final SdJobService sdJobService;
    private final SdImageProcessor sdImageProcessor;

    public SdJobProcess(SdJobService sdJobService, SdImageProcessor sdImageProcessor) {
        this.sdJobService = sdJobService;
        this.sdImageProcessor = sdImageProcessor;
    }

    @Async("getAsyncExecutor")
    public void checkTaskProgress(Long sdId, ApiKeyDto apiKeyDto) {
        while (handleProgress(sdId, apiKeyDto)) {
            try {
                Thread.sleep(1000);
            } catch (InterruptedException ignored) {
            }
        }
    }

    private boolean handleProgress(Long sdId, ApiKeyDto apiKeyDto) {
        SdJobEntity sdJobEntity = sdJobService.getById(sdId);
        if (sdJobEntity == null) {
            return false;
        }
        Integer nowProcess = sdJobEntity.getProgress();
        if (nowProcess == 100 || nowProcess == -1) {
            return false;
        }
        String apiURL = String.format("%s/sdapi/v1/progress?skip_current_image=false", apiKeyDto.getUrl());
        OkHttpClient okHttpClient = OkHttpClientUtil.getOkHttpClient(apiKeyDto.getProxyUrl(), true);
        Request request = new Request.Builder()
                .url(apiURL)
                .get()
                .build();
        Call call = okHttpClient.newCall(request);
        try (Response response = call.execute()) {
            if (response.isSuccessful() && response.body() != null) {
                String result = response.body().string();
                JSONObject jsonObject = JSONUtil.parseObj(result);
                int progress = (int) (jsonObject.getDouble("progress", 0.0) * 100);
                if (progress != 0 && progress > nowProcess) {
                    LambdaUpdateWrapper<SdJobEntity> updateWrapper = new LambdaUpdateWrapper<>();
                    updateWrapper.eq(SdJobEntity::getId, sdJobEntity.getId());
                    updateWrapper.set(SdJobEntity::getProgress, progress);
                    sdJobService.update(updateWrapper);
                }
                LOGGER.message("Check task progress").context("response", result).info();
            }
        } catch (IOException e) {
            LOGGER.message("SD图片请求失败").context("exception", e).error();
        }
        sdImageProcessor.notifyUpdateTask(sdJobEntity.getId());
        return true;
    }

}
